﻿using log4net;
using OpenCL.Net;
using SharpML.Computation.Extensions;
using SharpML.Computation.Enumerations;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using SharpML.Computation.Interfaces;
using System.Runtime.InteropServices;

namespace SharpML.Computation.Implementation
{
    /// <summary>
    /// Provides an engine which provides means to execute stub calculators.
    /// </summary>
    public class ComputationEngine : IDisposable
    {
        #region - Private Member Variables -

            // Simple logging pattern
            private static readonly ILog Log = LogManager.GetLogger( System.Reflection.MethodBase.GetCurrentMethod().DeclaringType );

            private static List<Device> CpuDevices;
            private static List<Device> GpuDevices;

            private static List<Context> CpuContexts;
            private static List<Context> GpuContexts;

            private static List<CommandQueue> CpuQueues;
            private static List<CommandQueue> GpuQueues;

            private static int RefCount;

            public static int IntPtrSize;

            private Context Context;
            private CommandQueue Queue;
            private Device Device;
            private List<IMem> Memory;
            private List<Kernel> Kernels;

        #endregion

        #region - OpenCL -

            public const string PsuedoRandomNumbers = @"
// Implementation based on George Marsaglia's MWC (multiply with carry) generator.
// George Marsaglia's MWC (multiply with carry) generator.
// Although it is very simple, it passes Marsaglia's DIEHARD series of random number generator tests.
// 
// Written by John D. Cook 
// http://www.johndcook.com
//
//
// Ported to OpenCL by Christopher J. Hanson
// Added feedback loop to interpolate randomization
// Call PsuedoRandom( &random_t ) for feedback

typedef struct { float Value; uint Seed_w; uint Seed_z; } random_t;

random_t Seed( uint seed )
{
    random_t result;
    result.Value = (float)seed;
    result.Seed_w = seed;
    result.Seed_z = seed;
    return result;
}

void NextUInt( random_t* random )
{
    uint m_w = random->Seed_w;
    uint m_z = random->Seed_z;
    random->Seed_w = 18000 * ( m_w & 65535 ) + ( m_w >> 16 );
    random->Seed_z = 36969 * ( m_z & 65535 ) + ( m_z >> 16 );
    random->Value = ( ( random->Seed_z << 16 ) + random->Seed_w );
}

void NextUniform( random_t* random )
{
    NextUInt( random );
    random->Value = ( random->Value + 1.0f ) * 2.328306435454494e-10;
}

void NextNormal( random_t* random )
{
    NextUniform( random );
    float v1 = random->Value;

    NextUniform( random );
    float v2 = random->Value;

    float r = sqrt( -2.0f * log( v1 ) );
    float theta = ( 2.0f * 3.1415926535897931 ) * v2;
    random->Value = r * sin( theta );

}

void PsuedoRandom( random_t* random, uint loops )
{
    loops %= 4;
    for( int index = 0; index < loops; index++ )
        NextUInt( random );
    NextUniform( random );
    loops = (uint)(random->Value * 10);
    loops %= 10;
    for( int index = 0; index < loops; index++ )
        NextUInt( random );
}
            ";
            
        #endregion

        #region - Construction -

            public ComputationEngine( ComputeDevices device, int ordinal )
            {
                // Contract
                if( ordinal > this[device] )
                    throw new ArgumentOutOfRangeException( "Device ordinal out of bounds" );

                // Localize
                var contexts = device == ComputeDevices.GPU ? GpuContexts : CpuContexts;
                var queues = device == ComputeDevices.GPU ? GpuQueues: CpuQueues;
                var devices = device == ComputeDevices.GPU ? GpuDevices : CpuDevices;

                // Select context and queue
                Context = contexts[ordinal];
                Queue = queues[ordinal];
                Device = devices[ordinal];

                // Track for statics disposal
                RefCount++;

                // Log it
                Log.InfoFormat( "Initialized Computation Engine {0}:{1}, reference {2}", device, ordinal, RefCount );
            }

            static ComputationEngine()
            {
                IntPtrSize = Marshal.SizeOf( typeof( IntPtr ) );
                ErrorCode error;
                var platforms = Cl.GetPlatformIDs(out error);
                if( error != ErrorCode.Success || platforms.Length == 0 )
                    throw new InvalidOperationException( string.Format( "Unable to retrieve an OpenCL Platform, error was: {0}!", error ) );

                GpuDevices = new List<Device>(Cl.GetDeviceIDs( platforms[0], DeviceType.Gpu, out error) );
                if( error != ErrorCode.Success || platforms.Length == 0 )
                    Log.Warn( "No GPU devices!" );

                CpuDevices = new List<Device>(Cl.GetDeviceIDs( platforms[0], DeviceType.Cpu, out error));
                if( error != ErrorCode.Success || platforms.Length == 0 )
                    Log.Warn( "No CPU devices?!" );

                GpuQueues = new List<CommandQueue>();
                CpuQueues = new List<CommandQueue>();

                GpuContexts = GatherContextsAndQueues(GpuDevices, GpuQueues);
                CpuContexts = GatherContextsAndQueues(CpuDevices, CpuQueues);
            }

        #endregion

        #region - Private Functionality -

            private static List<Context> GatherContextsAndQueues( List<Device> gpuDevices, List<CommandQueue> queues )
            {
                List<Context> results = new List<Context>();
                ErrorCode error;
                foreach( var device in gpuDevices )
                {
                    var context = Cl.CreateContext( null, 1, new[] { device }, null, IntPtr.Zero, out error );

                    results.Add( context );

                    if( error != ErrorCode.Success )
                        throw new InvalidOperationException( String.Format( "Unable to retrieve an OpenCL GPU Context, error was: {0}!", error ) );

                    queues.Add( Cl.CreateCommandQueue( context, device, (CommandQueueProperties)0, out error ) );

                    if( error != ErrorCode.Success )
                        throw new InvalidOperationException( String.Format( "Unable to retrieve an OpenCL GPU Queue, error was: {0}!", error ) );
                }
                return results;
            }

        #endregion

        #region - Public Functionality -


            public int this[ComputeDevices deviceType]
            {
                get { return deviceType == ComputeDevices.GPU ? GpuDevices.Count : CpuDevices.Count; }
            }

            public Kernel[] Compile( string[] sources, string[] kernels, IComputationMemory computeMemory )
            {
                if( Kernels != null || Kernels.Count != 0 )
                    throw new NotSupportedException( "You may only compile once per engine." );

                Kernels = new List<Kernel>();
                ErrorCode error;

                using( var program = Cl.CreateProgramWithSource( Context, 1, sources, null, out error ) )
                {
                    if( error != ErrorCode.Success )
                        // TODO: More descript error (per device)
                        throw new InvalidProgramException( String.Format( "Unable to run Cl.CreateProgramWithSource: {0}!", error ) );

                    error = Cl.BuildProgram( program, 1, new[] { Device }, string.Empty, null, IntPtr.Zero );

                    if( error != ErrorCode.Success )
                    {
                        char[] buffer = new char[90000];
                        InfoBuffer infoBuffer = new InfoBuffer( new IntPtr( 90000 ) );
                        IntPtr retSize;
                        Cl.GetProgramBuildInfo( program, Device, ProgramBuildInfo.Log, new IntPtr( 90000 ), infoBuffer, out retSize );

                        throw new InvalidProgramException( "Error building program:\n" + infoBuffer.ToString() );
                    }

                    var buildStatus = Cl.GetProgramBuildInfo( program, Device, ProgramBuildInfo.Status, out error ).CastTo<BuildStatus>();
                    if( buildStatus != BuildStatus.Success )
                        throw new InvalidProgramException( string.Format( "GetProgramBuildInfo returned {0} for program!", buildStatus ) );

                    foreach( string kernel in kernels )
                        Kernels.Add( program.CreateKernel( kernel ) );

                    computeMemory.Allocate( this );

                    return Kernels.ToArray();
                }
            }

            public IMem AllocateFloatArray( float[] array, MemFlags flags )
            {
                IMem result = Context.AllocateFloatArray( array, flags );
                Memory.Add( result );
                return result;
            }

            public IMem AllocateIntArray( int[] array, MemFlags flags )
            {
                IMem result = Context.AllocateIntArray( array, flags );
                Memory.Add( result );
                return result;
            }

            public void ReadFloatArray( IMem mem, float[] array, BlockModes blockMode )
            {
                mem.Read( Queue, sizeof( float ) * array.Length, array, blockMode );
            }

            public void WriteFloatArray( IMem mem, float[] array, BlockModes blockMode )
            {
                mem.Write( Queue, sizeof( float ) * array.Length, array, blockMode );
            }

            public void ReadIntArray( IMem mem, int[] array, BlockModes blockMode )
            {
                mem.Read( Queue, sizeof( int ) * array.Length, array, blockMode );
            }

            public void WriteIntArray( IMem mem, int[] array, BlockModes blockMode )
            {
                mem.Write( Queue, sizeof( int ) * array.Length, array, blockMode );
            }

            public void Enqueue( Kernel kernel, params int[] sizes )
            {
                Event clEvent;
                ErrorCode error = ErrorCode.Unknown;

                if( sizes != null )
                    error = Cl.EnqueueNDRangeKernel( Queue, kernel, (uint)sizes.Length, null, sizes.ToList().Select( size => new IntPtr( size ) ).ToArray(), null, 0, null, out clEvent );
                else
                    error = Cl.EnqueueTask( Queue, kernel, 0, null, out clEvent );

                if( error != ErrorCode.Success )
                    throw new InvalidOperationException( "EnqueueNDRangeKernel failed:" + error );

                clEvent.Dispose();
            }

     
            public int AvailableMemory
            {
                get
                {
                    ErrorCode error;
                    IntPtr size = Cl.GetDeviceInfo( Device, DeviceInfo.GlobalMemSize, out error ).CastTo<IntPtr>();
                    return size.ToInt32();
                }
            }

        #endregion

        #region - IDisposable Implementation -

            public void Dispose()
            {
                if( --RefCount == 0 )
                {
                    Log.Info( "Disposing OpenCL" );

                    foreach( var kernel in Kernels )
                        kernel.Dispose();

                    foreach( var memory in Memory )
                        memory.Dispose();
                    
                    foreach( var queue in GpuQueues.Union( CpuQueues ) )
                        queue.Dispose();

                    foreach( var context in GpuContexts.Union( CpuContexts ) )
                        context.Dispose();
                }
            }

        #endregion

            public void Flush()
            {
                Queue.Finish();
            }
    }
}
